from datasets import load_dataset
import pandas as pd
from collections import defaultdict
import re
import os


def count_chinese_chars(text):
    """统计中文字符数"""
    return len(re.findall(r'[\u4e00-\u9fff]', text))


def count_english_words(text):
    """统计英文单词数"""
    return len(text.split())


def save_to_excel(data, filename, max_rows=1000000):
    """将数据分批保存到多个Excel文件"""
    if not data:
        return

    # 创建输出目录
    os.makedirs('output3', exist_ok=True)

    # 计算需要多少个文件
    total_chunks = (len(data) + max_rows - 1) // max_rows

    for i in range(total_chunks):
        start_idx = i * max_rows
        end_idx = min((i + 1) * max_rows, len(data))
        chunk_data = data[start_idx:end_idx]

        # 生成文件名
        if total_chunks > 1:
            chunk_filename = f'output3/{filename}_{i + 1}.xlsx'
        else:
            chunk_filename = f'output3/{filename}.xlsx'

        # 保存数据
        df = pd.DataFrame(chunk_data)
        df.to_excel(chunk_filename, index=False)
        print(f"保存数据到 {chunk_filename}, 共 {len(chunk_data)} 条")


def classify_style_and_domain(text):
    """根据文本内容分类风格和领域"""
    text = text.lower()

    # 风格分类（放宽匹配条件）
    if re.search(r'(hereby|pursuant|whereas|shall|thereof|according|official|formally)', text):
        return '正式翻译', '请使用正式的语言风格将以下中文翻译成英文：'
    elif re.search(r'(haha|lol|wow|hey|cool|fun|awesome|nice|great)', text):
        return '非正式翻译', '请用轻松口语的方式将以下中文翻译成英文：'

    # 领域分类（扩大关键词范围）
    elif re.search(r'(tech|algorithm|software|database|computer|technology|system|data|digital|online|internet)', text):
        return '科技翻译', '请将以下技术相关的中文内容翻译成英文：'
    elif re.search(r'(market|sales|advertising|brand|customer|business|company|product|service)', text):
        return '营销翻译', '请将以下营销相关的中文内容翻译成英文：'
    elif re.search(r'(court|law|legal|regulation|pursuant|rights|justice|case)', text):
        return '法律翻译', '请将以下法律相关的中文内容翻译成英文：'
    elif re.search(r'(patient|medical|treatment|disease|doctor|health|hospital|clinic)', text):
        return '医疗翻译', '请将以下医疗相关的中文内容翻译成英文：'
    elif re.search(r'(student|teacher|education|school|learn|study|course|training)', text):
        return '教育翻译', '请将以下教育相关的中文内容翻译成英文：'
    else:
        return '基本翻译', '请将以下中文翻译成英文：'


# 加载数据集
print("正在加载数据集...")
ds = load_dataset("wmt/wmt19", "zh-en")
print(f"数据集大小: {len(ds['train'])} 条")

# 按不同类型收集数据
translation_data = defaultdict(list)
total_processed = 0
length_filtered = 0

# 字数限制
MIN_ZH_LENGTH = 300
MAX_ZH_LENGTH = 2000
MIN_EN_LENGTH = 20
MAX_EN_LENGTH = 3000

for example in ds['train']:
    total_processed += 1
    if total_processed % 10000 == 0:
        print(f"已处理 {total_processed} 条数据...")

    zh_text = example['translation']['zh']
    en_text = example['translation']['en']

    # 检查字数限制
    zh_length = count_chinese_chars(zh_text)
    en_length = count_english_words(en_text)

    if MIN_ZH_LENGTH <= zh_length <= MAX_ZH_LENGTH and MIN_EN_LENGTH <= en_length <= MAX_EN_LENGTH:
        length_filtered += 1
        # 分类并生成指令
        category, instruction = classify_style_and_domain(en_text)

        # 收集数据
        translation_data[category].append({
            '指令': instruction,
            '原文': zh_text,
            '译文': en_text
        })

print(f"\n符合字数要求的数据: {length_filtered} 条")
print("\n各类别数据统计:")
for category, data in translation_data.items():
    print(f"{category}: {len(data)}条")

# 分别保存每个类别的数据
for category, data in sorted(translation_data.items()):
    if category != '基本翻译' and len(data) > 0:
        save_to_excel(data, f'translation_dataset_{category}')

# 最后保存基本翻译
# if '基本翻译' in translation_data:
#     save_to_excel(translation_data['基本翻译'], 'translation_dataset_基本翻译')

print("\n所有数据已保存到 output2 目录")